#!/usr/bin/env python3
# coding: utf-8
# Copyright 2024 Huawei Technologies Co., Ltd
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
# ===========================================================================
import os
import platform
import re
import json
import tarfile
import glob

from ansible.module_utils.basic import AnsibleModule
from ansible.module_utils import common_info, common_utils, venv_installer
from ansible.module_utils.common_info import DeployStatus


class DifyInstaller:

    def __init__(self):
        self.module = AnsibleModule(
            argument_spec=dict(
                tags=dict(type="list", required=True),
                resources_dir=dict(type="str", required=True),
                dify_code_path=dict(type="str", required=True),
            )
        )
        self.run_tags = self.module.params["tags"]
        self.resources_dir = os.path.expanduser(self.module.params["resources_dir"])
        self.dify_code_path = os.path.expanduser(self.module.params["dify_code_path"])
        self.arch = platform.machine()
        self.messages = []

    def _module_failed(self):
        return self.module.fail_json(msg="\n".join(self.messages), rc=1, changed=False)

    def _module_success(self):
        return self.module.exit_json(msg="Install Dify success.", rc=0, changed=True)

    def _find_dify_pkgs(self, filename, pattern):
        pkgs, msgs = common_utils.find_files(os.path.join(self.resources_dir, filename), pattern)
        self.messages.extend(msgs)
        if not pkgs:
            if "auto" in self.run_tags:
                self.module.exit_json(std_out="[ASCEND] can not find Dify package, Dify install skipped", rc=0,
                                      result={DeployStatus.DEPLOY_STATUS: DeployStatus.SKIP},
                                      changed=False)
            else:
                self.messages.append("[ASCEND] can not find Dify package.")
                self._module_failed()

        return pkgs

    def load_dify_images(self):
        dify_image_pattern = "*{}.tar*".format(self.arch)
        dify_images_path = self._find_dify_pkgs("DifyImages*", dify_image_pattern)
        docker_path = self.module.get_bin_path("docker")
        if not docker_path:
            self.module.fail_json("[ASCEND][ERROR] Docker not installed, please install docker first.")

        # 获取本地已有的 repo:tag 列表
        rc, out, err = self.module.run_command([docker_path, "images", "--format", "{{{{.Repository}}}}:{{{{.Tag}}}}"])
        existing = set(out.splitlines()) if rc == 0 and out else set()

        image_names = []
        for image_path in dify_images_path:
            repo_tags = []
            # 尝试从 image tar 中读取 manifest.json 获取 RepoTags
            try:
                with tarfile.open(image_path, 'r') as tf:
                    try:
                        member = tf.extractfile('manifest.json')
                        if member:
                            manifest = json.load(member)
                            if isinstance(manifest, list) and manifest:
                                repo_tags = manifest[0].get('RepoTags', []) or []
                    except Exception:
                        repo_tags = []
            except Exception:
                repo_tags = []

            # 如果能解析到 repo_tags 并且本地已有所有 tag，则跳过加载
            if repo_tags and all(tag in existing for tag in repo_tags):
                self.messages.append("Image(s) {} already exist, skip loading {}".format(
                    ", ".join(repo_tags), image_path))
                image_names.extend(repo_tags)
                continue

            # 否则执行 docker load
            rc, out, err = self.module.run_command([docker_path, "load", "-i", image_path])
            if rc != 0:
                self.module.fail_json(
                    '[ASCEND][ERROR] Failed to loading Docker image from {}, error:{}'.format(image_path, err))
            self.messages.append("Loading Docker image from {} successfully.".format(image_path))
            for line in out.splitlines():
                if 'Loaded image:' in line:
                    image_names.append(line.split()[-1])

        return image_names

    def _check_net_ipv4_ip_forward(self):
        """检查并启用 IP 转发"""
        rc, out, err = self.module.run_command(["sysctl", "net.ipv4.ip_forward"])
        if rc != 0:
            self.module.fail_json("[ASCEND][ERROR] Failed to check net.ipv4.ip_forward: {}".format(err))
        if "net.ipv4.ip_forward = 1" in out:
            self.messages.append("[ASCEND] net.ipv4.ip_forward is already enabled.")
            return
        # 启用 IP 转发
        rc, out, err = self.module.run_command(["sysctl", "-w", "net.ipv4.ip_forward=1"])
        if rc != 0:
            self.module.fail_json("[ASCEND][ERROR] Failed to enable net.ipv4.ip_forward: {}".format(err))
        self.messages.append("[ASCEND] Enabled net.ipv4.ip_forward successfully.")

    def run_dify(self):
        # 检查 docker compose
        compose_v2 = False
        compose_path = self.module.get_bin_path("docker-compose", required=False)
        docker_path = self.module.get_bin_path("docker", required=False)
        if not compose_path:
            # 检查 docker compose v2
            rc, out, _ = self.module.run_command([docker_path, "compose", "version"])
            if rc != 0:
                self.module.fail_json("[ASCEND][ERROR] Docker Compose not installed, please install docker-compose or docker compose plugin.")
            compose_v2 = True
        
        # 解压 dify 代码
        code_pkgs = glob.glob(os.path.join(self.resources_dir, "DifyCodes*", "DifyCodes-*.tar.gz"))
        if not code_pkgs:
            self.module.fail_json("[ASCEND][ERROR] can not find DifyCodes in {}".format(os.path.join(self.resources_dir, "DifyCodes*")))
        code_pkg = code_pkgs[0]
        # 解压到dify_code_path
        #extract_dir = os.path.join(self.resources_dir, "DifyCodes")
        # 自动创建解压目的路径（关键步骤）
        # 使用 ansible module 的 file 模块创建目录，支持多级目录（recurse=yes）
        dir_result = self.module.run_command(
            ["mkdir", "-p", self.dify_code_path]  # mkdir -p：不存在则创建，存在则不报错（兼容多级目录）
        )
        # 检查目录创建结果
        if dir_result[0] != 0:
            self.module.fail_json('[ASCEND][ERROR] Failed to create target directory: {}. Error: {}'.format(
                self.dify_code_path, dir_result[2].strip()))
            return
        # 前置检查：tar 命令是否存在
        if not self.module.get_bin_path("tar"):
            self.module.fail_json('[ASCEND][ERROR] "tar" command is not installed, cannot extract code package.')
            return
        rc, out, err = self.module.run_command(["tar", "-xzf", code_pkg, "-C", self.dify_code_path])
        if rc != 0:
            self.module.fail_json("[ASCEND][ERROR] Failed to extract DifyCodes: {}".format(err))
        self.messages.append("Extracted DifyCodes successfully.")
        # 进入 Dify 源代码的 Docker 目录
        #dify_dirs = glob.glob(os.path.join(extract_dir, "dify*/"))
        dify_dirs = glob.glob(os.path.join(self.dify_code_path, "dify-*/"))
        if dify_dirs:
            dify_dir = dify_dirs[0].rstrip("/")
        else:
            self.module.fail_json(f"[ASCEND][ERROR] Dify source directory not found in {extract_dir}")
        docker_dir = os.path.join(dify_dir, "docker")
        if not os.path.isdir(docker_dir):
            self.module.fail_json(f"[ASCEND][ERROR] docker directory not found in {dify_dir}")
        # 加载镜像
        image_names = self.load_dify_images()
        # 复制环境配置文件
        env_pkgs = self._find_dify_pkgs("DifyCodes*", "dify_*.env.example")
        if not env_pkgs:
            env_example = os.path.join(docker_dir, ".env.example")
            self.messages.append("Default .env.example is bing used.")
        else:
            env_example = env_pkgs[0]
        env_file = os.path.join(docker_dir, ".env")
        # 强制覆盖已有的.env文件
        rc, _, err = self.module.run_command(["cp", env_example, env_file])
        if rc != 0:
            self.module.fail_json("[ASCEND][ERROR] Failed to copy .env.example to .env: {}".format(err))
        self.messages.append("Copied .env.example to .env.")
        # 检查并启用 IP 转发
        self._check_net_ipv4_ip_forward()
        # 启动 Docker 容器
        if compose_v2:
            cmd = [docker_path, "compose", "up", "-d"]
        else:
            cmd = [compose_path, "up", "-d"]
        rc, out, err = self.module.run_command(cmd, cwd=docker_dir)
        if rc != 0:
            self.module.fail_json("[ASCEND][ERROR] Failed to start Dify containers: {}".format(err))
        self.messages.append("Started Dify containers successfully.")
        # 检查容器状态
        self.check_dify_containers(image_names)

    def check_dify_containers(self, image_names):
        # 检查所有 image_names 的容器都为 up
        docker_path = self.module.get_bin_path("docker")
        rc, out, err = self.module.run_command([docker_path, "ps", "--format", "{{.Image}} {{.Status}}"])
        if rc != 0:
            self.module.fail_json("[ASCEND][ERROR] Failed to check container status: {}".format(err))
        running_images = {}
        for line in out.splitlines():
            parts = line.strip().split(None, 1)
            if len(parts) == 2:
                running_images[parts[0]] = parts[1]
        not_up = []
        for image in image_names:
            if image not in running_images or not running_images[image].lower().startswith("up"):
                not_up.append(image)
        if not_up:
            self.module.fail_json("[ASCEND][ERROR] The following Dify containers are not running: {}".format(", ".join(not_up)))
        self.messages.append("All Dify containers are running.")

    def run(self):
        try:
            self.run_dify()
        except Exception as e:
            self.messages.append(str(e))
            self._module_failed()
        self._module_success()


def main():
    DifyInstaller().run()


if __name__ == "__main__":
    main()
